# Copyright (c) 2018 Presto Labs Pte. Ltd.
# Author: jaewon

import logging

from coin.base.timestamp import get_timestamp
from coin.exchange.base.order_gateway import OrderSide, OrderType


class BitmexFastPullMMExecutor(object):
  def __init__(self, execution_params, order_gateway, logger=None):
    self._logger = logger or logging.getLogger('BitmexFastPullMMExecutor')
    self._og = order_gateway

    self._symbol = execution_params['product'].native_symbol
    self._lot_size = execution_params['lot_size']
    self._stickiness = execution_params['stickiness']
    self._min_pos = execution_params['min_pos']
    self._max_pos = execution_params['max_pos']
    self._order_update_period = (execution_params.get('order_update_period', 2.) * (10**9))
    self._post_only = execution_params.get('post_only', True)

    self._sell_on_the_fly_orders = []
    self._sell_working_orders = []
    self._buy_on_the_fly_orders = []
    self._buy_working_orders = []

    self._last_fire_time = get_timestamp()
    self._last_improvement_time = get_timestamp()
    self._create_list = []
    self._improving_amend_list = []
    self._pulling_amend_list = []

  def _update_orders(self):
    self._sell_on_the_fly_orders = []
    self._sell_working_orders = []
    self._buy_on_the_fly_orders = []
    self._buy_working_orders = []

    for order in self._og.get_working_order():
      if order.product != self._symbol:
        continue

      if order.order_side == OrderSide.SELL:
        if order.accepted:
          self._sell_working_orders.append(order)
        else:
          self._sell_on_the_fly_orders.append(order)

      elif order.order_side == OrderSide.BUY:
        if order.accepted:
          self._buy_working_orders.append(order)
        else:
          self._buy_on_the_fly_orders.append(order)

      else:
        self._logger.error('Unknown side: %s', str(order))
        continue

  def _is_order_amendable(self, order):
    if order.internal.cancel_sent:
      return False
    if order.internal.amend_sent:
      elapsed = get_timestamp() - order.internal.last_amend_time
      return elapsed > 1 * (10.0**9)
    return True

  def _cancel_sell_working_orders(self, exclude_order_id=None):
    for order in self._sell_working_orders:
      if order.order_id == exclude_order_id or \
         not self._is_order_amendable(order):
        continue
      self._pulling_amend_list.append(self._og.make_bitmex_amend(order.order_id, leaves_qty=0))

  def _cancel_buy_working_orders(self, exclude_order_id=None):
    for order in self._buy_working_orders:
      if order.order_id == exclude_order_id or \
         not self._is_order_amendable(order):
        continue
      self._pulling_amend_list.append(self._og.make_bitmex_amend(order.order_id, leaves_qty=0))

  def _manage_sell_order_cancel_only(self, sell_price):
    # If there is an on-the-fly order, cancel all working orders and keep the
    # on-the-fly order.
    if (sell_price is None or self._sell_on_the_fly_orders
        or self._og.get_position().get(self._symbol, 0) <= self._min_pos):
      self._cancel_sell_working_orders()
      return
    # If there is an working order that its price is same to sell_price,
    # keep it and cancel all other orders.
    exclude_order_id = None
    for order in self._sell_working_orders:
      if abs(order.price - sell_price) <= self._stickiness:
        exclude_order_id = order.order_id
        break
    self._cancel_sell_working_orders(exclude_order_id=exclude_order_id)

  def _manage_buy_order_cancel_only(self, buy_price):
    # If there is an on-the-fly order, cancel all working orders and keep the
    # on-the-fly order.
    if (buy_price is None or self._buy_on_the_fly_orders
        or self._og.get_position().get(self._symbol, 0) >= self._max_pos):
      self._cancel_buy_working_orders()
      return
    # If there is an working order that its price is same to buy_price,
    # keep it and cancel all other orders.
    exclude_order_id = None
    for order in self._buy_working_orders:
      if abs(order.price - buy_price) <= self._stickiness:
        exclude_order_id = order.order_id
        break
    self._cancel_buy_working_orders(exclude_order_id=exclude_order_id)

  def _manage_sell_order(self, sell_price, can_improve):
    # If there is an on-the-fly order, cancel all working orders and keep the
    # on-the-fly order.
    if (sell_price is None or self._sell_on_the_fly_orders
        or self._og.get_position().get(self._symbol, 0) <= self._min_pos):
      self._cancel_sell_working_orders()
      return

    # Create new order if there is no working order.
    if not self._sell_working_orders:
      if can_improve:
        order = self._og.make_bitmex_order(self._symbol,
                                           sell_price,
                                           self._lot_size,
                                           OrderSide.SELL,
                                           OrderType.LIMIT,
                                           post_only=self._post_only)
        self._create_list.append(order)
      return

    # If there is an working order that its price is same to sell_price,
    # keep it and cancel all other orders.
    for order in self._sell_working_orders:
      if abs(order.price - sell_price) <= self._stickiness:
        self._cancel_sell_working_orders(exclude_order_id=order.order_id)
        return

    # Choose one with the smallest price and modify it and
    # cancel all other orders.
    chosen_order = None
    for order in self._sell_working_orders:
      if self._is_order_amendable(order) and \
         (chosen_order is None or order.price < chosen_order.price):
        chosen_order = order

    if chosen_order:
      self._cancel_sell_working_orders(exclude_order_id=chosen_order.order_id)
      amend_obj = self._og.make_bitmex_amend(chosen_order.order_id,
                                             price=sell_price,
                                             leaves_qty=self._lot_size)
      if sell_price < chosen_order.price and can_improve:  # Improve
        self._improving_amend_list.append(amend_obj)
      elif sell_price > chosen_order.price:  # Pull
        self._pulling_amend_list.append(amend_obj)
    else:
      self._cancel_sell_working_orders()

  def _manage_buy_order(self, buy_price, can_improve):
    # If there is an on-the-fly order, cancel all working orders and keep the
    # on-the-fly order.
    if (buy_price is None or self._buy_on_the_fly_orders
        or self._og.get_position().get(self._symbol, 0) >= self._max_pos):
      self._cancel_buy_working_orders()
      return

    # Create new order if there is no working order.
    if not self._buy_working_orders:
      if can_improve:
        order = self._og.make_bitmex_order(self._symbol,
                                           buy_price,
                                           self._lot_size,
                                           OrderSide.BUY,
                                           OrderType.LIMIT,
                                           post_only=self._post_only)
        self._create_list.append(order)
      return

    # If there is an working order that its price is same to buy_price,
    # keep it and cancel all other orders.
    for order in self._buy_working_orders:
      if abs(order.price - buy_price) <= self._stickiness:
        self._cancel_buy_working_orders(exclude_order_id=order.order_id)
        return

    # Choose one with the biggest price and modify it and
    # cancel all other orders.
    chosen_order = None
    for order in self._buy_working_orders:
      if self._is_order_amendable(order) and \
         (chosen_order is None or order.price > chosen_order.price):
        chosen_order = order

    if chosen_order:
      self._cancel_buy_working_orders(exclude_order_id=chosen_order.order_id)
      amend_obj = self._og.make_bitmex_amend(chosen_order.order_id,
                                             price=buy_price,
                                             leaves_qty=self._lot_size)
      if buy_price > chosen_order.price and can_improve:  # Improve
        self._improving_amend_list.append(amend_obj)
      elif buy_price < chosen_order.price:  # Pull
        self._pulling_amend_list.append(amend_obj)
    else:
      self._cancel_buy_working_orders()

  def _is_enough_rate_limit(self):
    default_limit = 300
    lower_bound = 100
    window_size = 300 * (10.0**9)

    rate_limit = self._og.get_rate_limit() or default_limit
    remaining = self._og.get_rate_limit_remaining() or lower_bound
    estimated_remaining = min(
        rate_limit,
        remaining + (get_timestamp() - self._last_fire_time) * (rate_limit / window_size))

    return estimated_remaining > lower_bound and \
           self._og.get_rate_count() < remaining

  def manage_mm_orders(self, sell_price, buy_price):
    fired = False
    fired_improvement = False
    self._create_list = []
    self._improving_amend_list = []
    self._pulling_amend_list = []
    self._update_orders()

    if self._is_enough_rate_limit():
      can_improve_order = \
          get_timestamp() - self._last_improvement_time > \
          self._order_update_period
      self._manage_sell_order(sell_price, can_improve_order)
      self._manage_buy_order(buy_price, can_improve_order)
    else:
      self._manage_sell_order_cancel_only(sell_price)
      self._manage_buy_order_cancel_only(buy_price)

    if self._pulling_amend_list or self._improving_amend_list:
      amend_list = self._pulling_amend_list + self._improving_amend_list
      self._og.amend_bulk(amend_list)
      self._logger.debug('Amend %d orders', len(amend_list))
      fired = True
      if self._improving_amend_list:
        fired_improvement = True
    if self._create_list:
      self._og.submit_bulk(self._create_list)
      self._logger.debug('Create %d orders', len(self._create_list))
      fired = True
      fired_improvement = True

    if fired:
      self._logger.debug('Rate Limit %s / %s',
                         str(self._og.get_rate_limit_remaining()),
                         str(self._og.get_rate_limit()))
      self._last_fire_time = get_timestamp()
    if fired_improvement:
      self._last_improvement_time = get_timestamp()
    return fired
